# Code for log-gaussian radial basis function.

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Union, Tuple, List
from jax import nn as nn
from .rbf_base import RadialBasisFunctions, _rbf_register

@_rbf_register('log_gaussian')
class LogGaussianBasis(RadialBasisFunctions):
    r"""Log Gaussian type RBF.
    Args:
        r_max (Length):         Maximum distance. Defatul: 1 nm.
        r_min (Length):         Minimum distance. Default: 0.04 nm.
        sigma (float):          Simga. Default: 0.3 nm.
        delta (float):          Space interval. Default: 0.0512 nm.
        num_basis (int):        Number of basis functions. Defatul: None.
        rescale (bool):         Whether to rescale the output of RBF from -1 to 1. Default: True
        clip_distance (bool):   Whether to clip the value of distance. Default: False.
        r_ref (float):         Reference distance. Default: 1 nm.
        name (str):             Name of the module. Default: "gaussian_basis".
    """

    def __init__(self, 
                 fp_type = jnp.float32, 
                 r_max: float = 1, 
                 r_min: float = 0.04, 
                 sigma: float = 0.3,
                 delta: float = 0.0512,
                 num_basis: Optional[int] = None,
                 rescale: bool = True, 
                 clip_distance: bool = False,
                 r_ref: float = 1,
                 name: str = "log_gaussian_basis",
                 ):
        super().__init__(fp_type=fp_type,
                         r_max=r_max,
                         r_min=r_min,
                         num_basis=num_basis,
                         clip_distance=clip_distance,
                         name=name)

        if r_max <= r_min:
            raise ValueError("[utils/rbf/LogGaussianBasis] r_max should be larger than r_min.")
        
        self.r_range = self.r_max - self.r_min

        if num_basis is None and delta is None:
            raise TypeError('[utils/rbf/LogGaussianBasis] "num_basis" and "delta" cannot both be "None".')
        if num_basis is not None and num_basis <= 0:
            raise ValueError('[utils/rbf/LogGaussianBasis] "num_basis" must be larger than 0.')
        if delta is not None and delta <= 0:
            raise ValueError('[utils/rbf/LogGaussianBasis] "delta" must be larger than 0.')
        
        self.delta = delta
        self.sigma = sigma
        self.rescale = rescale
        self.r_ref = r_ref 

        ## Liyh: trouble here
        self.log_rmin = np.log(self.r_min/self.r_ref).astype(self.fp_type)
        self.log_rmax = np.log(self.r_max/self.r_ref).astype(self.fp_type)
        self.log_range = self.log_rmax - self.log_rmin
        if self.delta is None and self.num_basis is not None:
            self.offsets = jnp.linspace(self.log_rmin, self.log_rmax, self.num_basis, dtype=self.fp_type)
            self.delta = self.log_range / (self.num_basis - 1)
        else:
            if self.num_basis is None:
                self.num_basis = int(np.ceil(self.log_range / self.delta)) + 1
            self.offsets = self.log_rmin + jnp.arange(0, self.num_basis, dtype=self.fp_type) * self.delta
        
        self.coefficient = -0.5 * jnp.reciprocal(jnp.square(self.sigma))
        self.inv_ref = jnp.reciprocal(self.r_ref)

    def __call__(self, distance: jnp.ndarray) -> jnp.ndarray:
        """Compute gaussian type RBF.

        Args:
            distance (Array): Array of shape `(...)`. Data type is float.

        Returns:
            rbf (Array):      Array of shape `(..., K)`. Data type is float.

        """
        if self.clip_distance:
            distance = jnp.clip(distance, self.r_min, self.r_max)
        
        # (...,) -> (..., 1)
        log_r = jnp.log(distance * self.inv_ref) ## Liyh: The main difference between jax and ms
        log_r = jnp.expand_dims(log_r, axis=-1)
        # (..., 1) - (..., K) -> (..., K)
        log_diff = log_r - self.offsets
        rbf = jnp.exp(self.coefficient * jnp.square(log_diff))

        if self.rescale:
            rbf = rbf * 2.0 - 1.0
        
        return rbf
    
    def __str__(self):
        return 'LogGaussianBasis<>'
